import gym
from utils import *
from ch4.main import PolicyIteration,ValueIteration

env = gym.make("FrozenLake-v0")  # 创建环境
env = env.unwrapped  # 解封装才能访问状态转移矩阵P
env.render()  # 环境渲染，通常是弹窗显示或打印出可视化环境

def print_env_info():
    holes = set()
    ends = set()
    for s in env.P:
        for a in env.P[s]:
            for s_ in env.P[s][a]:
                if s_[2] == 1.0:  # 获得奖励为1，代表是目标
                    ends.add(s_[1])
                if s_[3] == True:
                    holes.add(s_[1])
    holes = holes - ends
    print("冰洞的索引：",holes)
    print("目标的索引：",ends)

    for a in env.P[14]:  # 查看目标左边一格的状态转移信息
        print(env.P[14][a])

if __name__ == '__main__':
    # print_env_info()
    action_meaning = ['<','v','>','^']
    theta = 1e-5
    gamma = 0.9
    # agent = PolicyIteration(env,theta,gamma)
    # agent.policy_iteration()
    # print_agent(agent,action_meaning,[5,7,11,12],[15])

    agent = ValueIteration(env,theta,gamma)
    agent.value_iteration()
    print_agent(agent,action_meaning,[5,7,11,12],[15])



